#!/opt/local/bin/python

from scipy import pi, log, arctan, exp, cos, sin, inf, nan
from numpy import array, arange, complex128
import time

def print_timing(func):
    def wrapper(*args, **kwargs):
        t1 = time.time()
        res = func(*args, **kwargs)
        t2 = time.time()
        print '%s took %0.3f ms' % (func.func_name, (t2-t1)*1000.0)
        return res
    return wrapper

class transientcavity(object):
    def __init__(self, freq=704.42e6, QL=8e5, RoverQ=477.0, detuning=0):
        self.freq = freq
        self.detuning = detuning
        self.QL = QL
        self.RoverQ = RoverQ

        self.calcParams()
        self.calcMats()

    def calcParams(self):
        self.omega0    = 2*pi * self.freq
        self.omegahalf = self.omega0 / (2*self.QL)
        self.delomega  = 2*pi * self.detuning

        self.RL = self.RoverQ * self.QL

        self.phi = arctan (self.QL * ((self.omega0/(self.omega0-self.delomega)) - ((self.omega0-self.delomega)/self.omega0)) )
        self.cmplxRL = (1 + exp(2*self.phi*1j)) * self.RL
        
        self.tau = 1 / self.omegahalf

    def calcMats(self):
        self.Acmplx = -self.omegahalf + self.delomega*1j
        self.Bcmplx = self.RL * self.omegahalf

    def phasePred(self, t):
        omegahalf = self.omegahalf
        delomega = self.delomega

        VpredReal = exp(-omegahalf*t) * (-omegahalf*cos(delomega*t) + delomega*sin(delomega*t)) + omegahalf
        VpredImag = exp(-omegahalf*t) * (-omegahalf*sin(delomega*t) - delomega*cos(delomega*t)) + delomega

        return arctan(VpredImag / VpredReal) * 180/pi

class analysis(object):
    def __init__(self, tcavobj, Ib=50e-3, timestop=4e-3, acc=1e-3):
        self.timestep      = log(2) * tcavobj.tau * acc
        print "Timestep = %0.4e s" % self.timestep
        self.timestop      = timestop
        self.time          = arange(0, timestop, self.timestep)
        self.Acmplx        = tcavobj.Acmplx
        self.Bcmplx        = tcavobj.Bcmplx
        self.cmplxRl       = tcavobj.cmplxRL
        self.Acmplxtstep   = self.timestep * self.Acmplx
        self.Bcmplxtstep   = self.timestep * self.Bcmplx
        self.Acmplxtplus1  = self.Acmplxtstep + 1
        self.setpt         = nan
        self.dofb          = False
        self.Ibx           = Ib * (1 + 0j)

    def closeLoop(self):
        self.dofb = True

    def openLoop(self):
        self.dofb = False

    def feedbackConds(self, cmplxV, gain, tinj=(0, inf)):
        self.fbtime  = tinj
        self.feedback = feedback(setpt=cmplxV, gain=gain, tinj=tinj, cmplxRl=self.cmplxRl)

    @print_timing
    def run(self, iniV, Ig, Ibeam):
        I = Ig + Ibeam
        self.iniVcmplx = complex(iniV)

        FBcorrection = complex128(0 + 0j)
        self.feedback.fbsig = [FBcorrection] * len(self.time)

        Vcmplx = array([0 + 0j] * len(self.time))
        Vcmplx[0] = self.iniVcmplx

        Acmplxtplus1 = self.Acmplxtplus1
        Bcmplxtstep = self.Bcmplxtstep
        step = self.feedback.step
        dofb = self.dofb
        fbtime = self.fbtime
        time = self.time
        Ibx = self.Ibx
        for counter in xrange(len(self.time[1:])):
            Vcmplx[counter+1] = Acmplxtplus1 * Vcmplx[counter] + (Bcmplxtstep * (I[counter] + FBcorrection))

            if dofb and fbtime[0] < time[counter] < fbtime[1]:
                FBcorrection = step(Vcav=Vcmplx[counter+1], Ibx=Ibx, Ig=Ig[counter], counter=counter)
            elif dofb and time[counter] > fbtime[1]:
                FBcorrection = 0 + 0j
                self.feedback.fbsig[counter] = FBcorrection
            elif dofb and time[counter] < fbtime[0]:
                self.feedback.fbsig[counter] = FBcorrection
        self.V = Vcmplx

def calcCurrents(cavobj, tcavobj, time, pulselength, RoverQscaling=1.0, beamerr=(0.0, 0.0)):
    #beamgen = abs(cavobj.I_b)*cavobj.RoverQ*RoverQscaling*cavobj.Ql*cos(tcavobj.phi)
    beamgen = abs(cavobj.I_b0)*cavobj.RoverQ*RoverQscaling*cavobj.Ql*cos(tcavobj.phi)
    if beamgen==0:
        beamgen = 1e-15
    finalV  = abs(cavobj.V_cav) + beamgen
    tinj = tcavobj.tau * log(finalV / beamgen)
    
    drivecur = 2 * (cavobj.V_g / tcavobj.cmplxRL)
    predrivecur = 2 * (cavobj.V_g0 / tcavobj.cmplxRL)
    predrivephase = -(tcavobj.phasePred(tinj) - cavobj.beamph) * pi/180
    predrive = exp(predrivephase*1j) * (abs(predrivecur) + 0j)
    beamcur  = -abs(cavobj.I_b) * (1 + 0j) * (1+beamerr[0]) * exp(1j * (beamerr[1]*pi/180))
    
    Ibeam, Ig = [0 + 0j] * len(time), [0 + 0j] * len(time)
    Ibeam, Ig = array(Ibeam), array(Ig)
    for tstep in xrange(len(time)):
        if time[tstep] < tinj:
            Ig[tstep]    = predrive
        elif time[tstep] < (tinj + pulselength):
            Ibeam[tstep] = beamcur
            Ig[tstep]    = drivecur

    return (Ig, Ibeam, tinj)

class feedback(object):
    def __init__(self, setpt, gain, cmplxRl, tinj=(0, inf)):
        self.setpt   = setpt
        self.Pgain   = gain[0]
        self.Igain   = gain[1]
        self.Icorr   = 0 + 0j
        self.tinj    = tinj
        self.cmplxRl = complex(cmplxRl)
        self.count   = 0
        self.fbsig   = []

    def step(self, Vcav, Ibx, Ig, counter):
        Igret = (2 * Vcav / self.cmplxRl) + Ibx
        diff = Ig - Igret
        self.Icorr += 0.999 * diff * self.Igain
        self.count += 1
        self.fbsig[counter] = complex128((diff * self.Pgain) + self.Icorr)
        return self.fbsig[counter]

